package j2EE.servlet;


import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;

/**
 * Created with Intel IDEA.
 *
 * @Auther: gjt
 * @Date: 2024/11/30 14:39
 */
public class CacheResponseWrapper extends HttpServletResponseWrapper {
    // 替换OutputStream
    private CacheOutputStream cacheOutputStream;
    // 替换wrapper
    private PrintWriter printWriter;

    // servletOutputStream 的一个简单实现
    // 它保存了向其中写的数据
    class CacheOutputStream extends ServletOutputStream {
        private ByteArrayOutputStream bos;

        public CacheOutputStream() {
            bos = new ByteArrayOutputStream();
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener writeListener) {

        }

        @Override
        public void write(int b) throws IOException {
            bos.write(b);
        }

        // 读回存储数据
        protected byte[] getBytes() {
            return bos.toByteArray();
        }
    }

    public CacheResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    // 重写ServletResponseWrapper
    public ServletOutputStream getOutputStream() throws IOException {
        if (null != cacheOutputStream) return cacheOutputStream;

        // 确保只有一个OutputStream 或者 一个Writer
        if (null != printWriter) throw new IOException("Writer already in use");

        cacheOutputStream = new CacheOutputStream();
        return cacheOutputStream;
    }

    public PrintWriter getPrintWriter() throws IOException{
        if(null != printWriter)return printWriter;

        if(null != cacheOutputStream)throw new IOException("OutputStream already in use");

        printWriter = new PrintWriter(new OutputStreamWriter(new CacheOutputStream()));
        return printWriter;
    }

    protected byte[] getByte(){
        if(null == cacheOutputStream)return null;
        return cacheOutputStream.getBytes();
    }
}
